#include"fast_lio/imu_handle.h"

namespace fast_lio{

void IMUHandle::IMUInit(shared_ptr<Frame> p_frame){
    V3D cur_acc, cur_gyr,mean_acc,mean_gyr;
    int N = 1;
    mean_acc << p_frame->imu_deq_.front()->linear_acceleration.x,
                            p_frame->imu_deq_.front()->linear_acceleration.y, p_frame->imu_deq_.front()->linear_acceleration.z; //记录最早加速度为加速度均值
    mean_gyr <<p_frame->imu_deq_.front()->angular_velocity.x,
                            p_frame->imu_deq_.front()->angular_velocity.y, p_frame->imu_deq_.front()->angular_velocity.z; //记录最早角速度为角速度均值
  
    //读取所有imu信息
    for (const auto &imu : p_frame->imu_deq_)
    {
        cur_acc << imu->linear_acceleration.x, imu->linear_acceleration.y, imu->linear_acceleration.z;//记录当前加速度
        cur_gyr << imu->angular_velocity.x, imu->angular_velocity.y, imu->angular_velocity.z;//记录当前角速度
        mean_acc      += (cur_acc - mean_acc) / N;//更新加速度均值
        mean_gyr      += (cur_gyr - mean_gyr) / N;//更新角速度均值
        // cov_acc_ = cov_acc_ * (N - 1.0) / N + (cur_acc - mean_acc).cwiseProduct(cur_acc - mean_acc) * (N - 1.0) / (N * N); //更新加速度协方差（向量）
        // cov_gyr_ = cov_gyr_ * (N - 1.0) / N + (cur_gyr - mean_gyr).cwiseProduct(cur_gyr - mean_gyr) * (N - 1.0) / (N * N); //更新角速度协方差（向量）
        N ++;
    }
    p_frame->p_state_->gra_ = -mean_acc/mean_acc.norm()*G0_;
    // p_frame->p_state_->bg_ = mean_gyr;
    p_imu_last_ = p_frame->imu_deq_.back();

    init_acc_ = mean_acc;
    // addIMUPose(p_frame->getState(),p_imu_last_);
}

void IMUHandle::predictBatch(shared_ptr<Frame> p_frame){
    double dt = 0;
    V3D angvel_avr,acc_avr;
    Eigen::Matrix<double,State::DIM_OF_STATE_,State::DIM_OF_STATE_> Q;
    Q.setZero();

    acc_avr   <<0.5 * (p_imu_last_->linear_acceleration.x + p_frame->imu_deq_.front()->linear_acceleration.x),
        0.5 * (p_imu_last_->linear_acceleration.y +  p_frame->imu_deq_.front()->linear_acceleration.y),
        0.5 * (p_imu_last_->linear_acceleration.z +  p_frame->imu_deq_.front()->linear_acceleration.z);
    angvel_avr<<0.5 * (p_imu_last_->angular_velocity.x +  p_frame->imu_deq_.front()->angular_velocity.x),
                0.5 * (p_imu_last_->angular_velocity.y +  p_frame->imu_deq_.front()->angular_velocity.y),
                0.5 * (p_imu_last_->angular_velocity.z +  p_frame->imu_deq_.front()->angular_velocity.z);
    dt = p_frame->imu_deq_.front()->header.stamp.toSec()-p_imu_last_->header.stamp.toSec();
    predict(p_frame,acc_avr,angvel_avr,dt, Q);
    addIMUPose(p_frame->getState(),p_frame->imu_deq_.front());

    //imu debug
    // for(auto p:p_frame->imu_deq_)
    //     cout<<"imu time "<<p->header.stamp.toSec()<<endl;

    for (auto it_imu = p_frame->imu_deq_.begin(); it_imu < (p_frame->imu_deq_.end() - 1); it_imu++)
    {
        auto &&head = *(it_imu);
        auto &&tail = *(it_imu + 1);
        //加速度和角速度均值，j与j+1的均值
        acc_avr   <<0.5 * (head->linear_acceleration.x + tail->linear_acceleration.x),
            0.5 * (head->linear_acceleration.y + tail->linear_acceleration.y),
            0.5 * (head->linear_acceleration.z + tail->linear_acceleration.z);
        angvel_avr<<0.5 * (head->angular_velocity.x + tail->angular_velocity.x),
                    0.5 * (head->angular_velocity.y + tail->angular_velocity.y),
                    0.5 * (head->angular_velocity.z + tail->angular_velocity.z);
        
        acc_avr = acc_avr*G0_/init_acc_.norm();

        dt = tail->header.stamp.toSec() - head->header.stamp.toSec();//相邻imu时刻时间差

        Q.block<3, 3>(3, 3).diagonal() = dt*dt*cov_acc_;
        Q.block<3, 3>(6, 6).diagonal() = dt*dt*cov_gyr_;
        Q.block<3, 3>(9, 9).diagonal() = dt*cov_bias_gyr_;
        Q.block<3, 3>(12, 12).diagonal() = dt*cov_bias_acc_;
        //向前传播
        predict(p_frame,acc_avr,angvel_avr,dt, Q);//根据输入数据向前传播(两imu数据间隔时间差，协方差矩阵，输入数据） todo

        addIMUPose(p_frame->getState(),tail);
    }
    p_imu_last_ = p_frame->imu_deq_.back();

    undistortPoints(p_frame);
}

void IMUHandle::predict(shared_ptr<Frame> p_frame,V3D& acc_avr,V3D& angvel_avr,double dt,
    Eigen::Matrix<double,State::DIM_OF_STATE_,State::DIM_OF_STATE_>& Q){
            // Average acceleration and angular rate
    State state_tmp; 
    state_tmp = p_frame->getState();

    MXD Ft =MXD::Identity(State::DIM_OF_STATE_, State::DIM_OF_STATE_);
    Ft.block<3, 3>(State::IDX_pos_, State::IDX_vel_) = M3D::Identity()*dt;

    Ft.block<3, 3>(State::IDX_vel_, State::IDX_q_) = 
        -state_tmp.q_.toRotationMatrix() * Sophus::SO3d::hat(acc_avr - state_tmp.ba_)*dt;
    Ft.block<3, 3>(State::IDX_vel_, State::IDX_ba_) =
        -state_tmp.q_.toRotationMatrix()*dt;
    Ft.block<3, 3>(State::IDX_vel_, State::IDX_gra_) = M3D::Identity()*dt;
    // Eigen::Matrix<double, 2, 1> vec = Eigen::Matrix<double, 2, 1>::Zero();
	// Eigen::Matrix<double, 3, 2> grav_matrix;
	// state_tmp.gra_.S2_Mx(vec,grav_matrix);
    // Ft.block<3, 2>(State::IDX_vel_, State::IDX_gra_) = grav_matrix*dt;

    Ft.block<3, 3>(State::IDX_q_, State::IDX_q_) =
       Sophus::SO3d::exp(-(angvel_avr-state_tmp.bg_)*dt).matrix();
    Ft.block<3, 3>(State::IDX_q_, State::IDX_bg_) = -M3D::Identity()*dt;

    // jacobian_ = F * jacobian_;
    state_tmp.cov_P_ =
        Ft * state_tmp.cov_P_ * Ft.transpose() + Q;
    
    V3D un_acc = state_tmp.q_.toRotationMatrix() * (acc_avr - state_tmp.ba_) + state_tmp.gra_;
    // V3D un_acc = state_tmp.q_.toRotationMatrix() * (acc_avr - state_tmp.ba_) + V3D(state_tmp.gra_[0],state_tmp.gra_[1],state_tmp.gra_[2]);
    state_tmp.pos_ = state_tmp.pos_+state_tmp.vel_*dt; //+0.5*un_acc*dt*dt;
    state_tmp.vel_ = state_tmp.vel_+un_acc*dt;
    V3D un_gyr = angvel_avr - state_tmp.bg_;
    Q4D dq(Sophus::SO3d::exp(un_gyr * dt).matrix());
    state_tmp.q_ = state_tmp.q_ * dq;
    state_tmp.q_ = state_tmp.q_.normalized();
    p_frame->setState(state_tmp);
}

void IMUHandle::undistortPoints(shared_ptr<Frame> p_frame){
    int idx_pt = p_frame->pc_full_->points.size()-1;
    double dt;
    Eigen::Matrix4d T_ref_w = Eigen::Matrix4d::Identity();
    T_ref_w.block<3, 3>(0, 0) = IMU_pose_v_.back().q.toRotationMatrix();
    T_ref_w.block<3, 1>(0, 3) = IMU_pose_v_.back().pos;
    T_ref_w = T_ref_w.inverse().eval();
    for(int idx_imu=IMU_pose_v_.size()-1;idx_imu>=0;idx_imu--){
        for(;idx_pt>=0&&
            p_frame->pc_full_->points[idx_pt].curvature+p_frame->time_stamp_>=IMU_pose_v_[idx_imu].time;idx_pt--){
            dt = p_frame->pc_full_->points[idx_pt].curvature+p_frame->time_stamp_ - IMU_pose_v_[idx_imu].time;
            V3D pos = IMU_pose_v_[idx_imu].pos + IMU_pose_v_[idx_imu].vel*dt+0.5*IMU_pose_v_[idx_imu].un_acc*dt*dt;
            Q4D q(Sophus::SO3d::exp(IMU_pose_v_[idx_imu].un_gyr * dt).matrix());
            q = IMU_pose_v_[idx_imu].q * q;
            Eigen::Matrix4d T1 = Eigen::Matrix4d::Identity();
            T1.block<3, 3>(0, 0) = q.toRotationMatrix();
            T1.block<3, 1>(0, 3) = pos;
            T1 = T_ref_w*T1;
            tranformPoint(p_frame->pc_full_->points[idx_pt],p_frame->pc_full_->points[idx_pt],T1);
        }
    }
    //test
    p_frame->pc_surf_.reset(new pcl::PointCloud<PointType>);
    pcl::copyPointCloud(*p_frame->pc_full_,*p_frame->pc_surf_);
    pcDownSample(p_frame->pc_surf_, 0.5);

    IMUPose imu_pose = IMU_pose_v_.back();
    IMU_pose_v_.clear();
    IMU_pose_v_.push_back(imu_pose);
}

void IMUHandle::addIMUPose(State state,sensor_msgs::ImuConstPtr imu_p){
    IMUPose imu_pose;
    imu_pose.un_acc = state.q_.toRotationMatrix() * (
        V3D(imu_p->linear_acceleration.x,imu_p->linear_acceleration.y,imu_p->linear_acceleration.z) - state.ba_) 
        + state.gra_;
    imu_pose.un_gyr = V3D(imu_p->angular_velocity.x,imu_p->angular_velocity.y,imu_p->angular_velocity.z) - state.bg_;
    imu_pose.pos = state.pos_;
    imu_pose.vel = state.vel_;
    imu_pose.q = state.q_;
    imu_pose.time = imu_p->header.stamp.toSec();
    IMU_pose_v_.push_back(imu_pose);
}

}
